from langchain_openai import ChatOpenAI
from langchain.prompts import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
import json
from utils.utils import fix_json_string
from utils.log_manager import app_logger

class ExamGenerator:
    
    def __init__(self,api_key,api_base,model_name="deepseek-chat"):
        self.llm = ChatOpenAI(
            temperature=0.9,
            model=model_name,
            openai_api_key=api_key,
            openai_api_base=api_base
        )
        
        # 定义输出格式
        self.response_schemas = [
            ResponseSchema(name="questions", description="生成的试题数组")
        ]
        self.output_parser = StructuredOutputParser.from_response_schemas(self.response_schemas)
        

    def set_prompt(self,num_questions,question_types,difficulty):
        self.system_template = """
            你是一名专业的出题老师，根据用户输入内容,请生成{num_questions}道、难度为{difficulty}、类型为{question_types}的试题。
            如果存在多个类型的试题,请遵循以下比例分配,否则按照100%比例生成指定类型的试题:
            - 60% 单选题 
            - 20% 多选题
            - 20% 简答题
            当前试题难度为{difficulty}
            当前的试题类型为{question_types}
            当前要求生成的试题数量为{num_questions}

            试题要求:
            - single:每题只有一个正确答案,提供4个选项(A, B, C, D)。
            - multiple:每题有多个正确答案,提供至少4个选项(A, B, C, D),并指明正确答案。  
            - fillBlank:提出一个问题,需要用户根据问题给出答案。
            - comprehensive: 综合题，包含single、multiple、fillBlank, 按照比例生成。

            请确保:
            1. 难度定义:
                - easy:测试基础知识,问题直白,选项较易辨认。
                - medium:测试理解和应用能力,选项有一定的干扰性。
                - hard:测试深入理解和分析能力,选项更具迷惑性。
            2. 生成的题目和选项要确保具有逻辑性,不包含歧义。
            当前试题难度为{difficulty}
            当前的试题类型为{question_types}
            当前要求生成的试题数量为{num_questions}
            
           ### 输出格式 
           请严格按照如下格式仅输出JSON，不要输出代码或其他信息
           options和correctAnswer生成的内容中不要包含A、B、C、D等选项，直接生成选项内容：
           {{
                "questions": [
                    {{
                    "id": 1,
                    "type": "single",
                    "text": "这是问题",
                    "options": ["A内容", "B内容", "C内容", "D内容"],
                    "correctAnswer": ["A内容"],
                    "explanation": "简要的解释"
                    }},
                    {{
                    "id": 2,
                    "type": "multiple",
                    "text": "这是问题",
                    "options": ["A内容", "B内容", "C内容", "D内容"],
                    "correctAnswer": ["A内容", "B内容"],
                    "explanation": "简要的解释"
                    }},
                    {{
                    "id": 3,
                    "type": "fillBlank",
                    "text": "这是问题",
                    "correctAnswer": ["答案"],
                    "explanation": "简要的解释"
                    }}
                ]
            }}
        """
        self.prompt = ChatPromptTemplate(
            messages=[
                SystemMessagePromptTemplate.from_template(
                    self.system_template,
                    partial_variables={
                       "num_questions":num_questions,
                       "difficulty":difficulty,
                       "question_types":question_types
                    }
                ),
                # 缓存占位符
                # MessagesPlaceholder(variable_name="chat_history"),
                HumanMessagePromptTemplate.from_template("参考内容：{reference_content}")
            ]
        )
        self.conversation = self.prompt | self.llm
        
    def generate_exam(self, reference_content):
        """生成试题"""
        try:
            # 调用模型生成结果
            result = self.conversation.invoke({
                "reference_content": reference_content,
            })
            app_logger.info(f"AGENT API Response: {result.content}")
            
            if not result.content:
                return []
            
            # 尝试解析结果
            return self._parse_result(result.content)
            
        except Exception as e:
            app_logger.error(f"试题生成过程中发生错误: {e}")
            return []
        
    def _parse_result(self, content):
        """解析LLM返回的结果"""
        # 修复常见的JSON格式错误
        fixed_content = fix_json_string(content)
        
        # 尝试使用输出解析器解析
        try:
            return self.output_parser.parse(fixed_content)
        except Exception as parser_error:
            app_logger.error(f"输出解析器失败: {parser_error}")
        
        # 尝试提取并解析JSON
        try:
            import re
            json_match = re.search(r'```json\s*([\s\S]*?)\s*```|({[\s\S]*})', fixed_content)
            if json_match:
                json_str = json_match.group(1) or json_match.group(2)
                parsed_data = json.loads(json_str)
                # 确保返回格式正确
                if "questions" in parsed_data:
                    return parsed_data
                else:
                    return {"questions": parsed_data}
        except Exception as json_error:
            app_logger.error(f"JSON手动解析失败: {json_error}")
        
        # 所有解析方法都失败，返回空列表
        return []
  

if __name__ == "__main__":
    api_key = "979e8470f5874c72883fda020d6b6222.eN1TVXQxZehaQxt2"
    api_base = "https://open.bigmodel.cn/api/paas/v4/"
    
    generator = ExamGenerator(api_key, api_base,model_name="glm-4")
    
    with open("E:\\APP\\QuickLearn\\server\\static\\docs\\test.txt", "r", encoding="utf-8") as file:
        reference_content = file.read()
    num_questions = 10
    question_types = "单选题、简答题"
    
    # 测试不同难度的试题生成
    difficulty = "easy"
    generator.set_prompt(num_questions,question_types,difficulty)
    exam = generator.generate_exam(reference_content)
    print(json.dumps(exam, ensure_ascii=False, indent=2))
